import os

from PIL import Image
from utils.logger import logger


class ImageProcessor:
    def __init__(self):
        # 使用BLIP模型作为替代方案
        try:
            from transformers import BlipForConditionalGeneration, BlipProcessor

            self.processor = BlipProcessor.from_pretrained(
                "Salesforce/blip-image-captioning-base"
            )
            self.model = BlipForConditionalGeneration.from_pretrained(
                "Salesforce/blip-image-captioning-base"
            )
        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}")
            raise

    def generate_text_from_image(self, image_path):
        """
        根据图片生成文字描述
        :param image_path: 图片文件路径
        :return: 生成的文字描述
        """
        # 验证图片文件存在
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图片文件不存在: {image_path}")

        # 加载并处理图片
        image = Image.open(image_path)
        inputs = self.processor(images=image, return_tensors="pt")

        # 生成文字描述
        generated_ids = self.model.generate(**inputs)
        generated_text = self.processor.batch_decode(
            generated_ids, skip_special_tokens=True
        )[0]

        return generated_text
